# -*- coding: utf-8 -*-
# @Time    : 2020/6/18 上午12:28
# @Author  : caotian
# @FileName: tbpaddletrain.py
# @Software: PyCharm
import paddle
import paddle.fluid as fluid
from paddle.fluid.dygraph.nn import Conv2D,Pool2D,Linear
import json
import gzip
import numpy as np
import random
from PIL import Image
import matplotlib.pyplot as plt
import os
import sys
curpath=os.path.abspath(os.curdir)
sys.path.append(curpath)
import optimizationdata as od
import optimizationmodel as om
from tb_paddle import SummaryWriter

data_writer=SummaryWriter(logdir="data")
with fluid.dygraph.guard():
    model=om.MNIST()
    model.train()
    train_loader=od.load_data('train')
    # 四种优化算法的设置方案，可以逐一尝试效果
    optimizer = fluid.optimizer.SGDOptimizer(learning_rate=0.01, parameter_list=model.parameters())

    EPOCH_NUM = 10
    iter = 0
    for epoch_id in range(EPOCH_NUM):
        for batch_id, data in enumerate(train_loader()):
            # 准备数据，变得更加简洁
            image_data, label_data = data
            image = fluid.dygraph.to_variable(image_data)
            label = fluid.dygraph.to_variable(label_data)

            # 前向计算的过程，同时拿到模型输出值和分类准确率
            predict, avg_acc = model(image, label)

            # 计算损失，取一个批次样本损失的平均值
            loss = fluid.layers.cross_entropy(predict, label)
            avg_loss = fluid.layers.mean(loss)

            # 每训练了100批次的数据，打印下当前Loss的情况
            if batch_id % 100 == 0:
                print("epoch: {}, batch: {}, loss is: {}, acc is {}".format(epoch_id, batch_id, avg_loss.numpy(),avg_acc.numpy()))
                data_writer.add_scalar('train/loss',avg_loss.numpy(),iter)
                data_writer.add_scalar('train/accuracy',avg_acc.numpy(),iter)
                iter=iter+100
            avg_loss.backward()
            optimizer.minimize(avg_loss)
            model.clear_gradients()
    fluid.save_dygraph(model.state_dict(),'mnist-model')

# 步骤3：命令行启动 tensorboard。
# 使用“tensorboard --logdir [数据文件所在文件夹路径] 的命令启动Tensor board。在Tensor board启动后，命令行会打印出可用浏览器查阅图形结果的网址。
#
# $ tensorboard --logdir log/data
# 步骤4：打开浏览器，查看作图结果，如 图6 所示。
# 查阅的网址在第三步的启动命令后会打印出来（如TensorBoard 2.0.0 at http://localhost:6006/），将该网址输入浏览器地址栏刷新页面的效果如下图所示。除了右侧对数据点的作图外，左侧还有一个控制板，可以调整诸多作图的细节。


